/*
 * Copyright 2011-2020 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package io.lettuce.core;

import static io.lettuce.core.ConnectionEventTrigger.local;
import static io.lettuce.core.ConnectionEventTrigger.remote;
import static io.lettuce.core.PlainChannelInitializer.pingBeforeActivate;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.security.GeneralSecurityException;
import java.time.Duration;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.function.Supplier;

import javax.net.ssl.SSLEngine;
import javax.net.ssl.SSLException;
import javax.net.ssl.SSLHandshakeException;
import javax.net.ssl.SSLParameters;

import io.lettuce.core.event.connection.ConnectedEvent;
import io.lettuce.core.event.connection.ConnectionActivatedEvent;
import io.lettuce.core.event.connection.DisconnectedEvent;
import io.lettuce.core.internal.HostAndPort;
import io.lettuce.core.internal.LettuceAssert;
import io.lettuce.core.protocol.AsyncCommand;
import io.lettuce.core.resource.ClientResources;
import io.netty.buffer.ByteBufAllocator;
import io.netty.channel.Channel;
import io.netty.channel.ChannelDuplexHandler;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.ssl.SslContext;
import io.netty.handler.ssl.SslContextBuilder;
import io.netty.handler.ssl.SslHandler;
import io.netty.handler.ssl.SslHandshakeCompletionEvent;
import io.netty.handler.ssl.util.InsecureTrustManagerFactory;

/**
 * Connection builder for SSL connections. This class is part of the internal API.
 *
 * @author Mark Paluch
 * @author Amin Mohtashami
 */
public class SslConnectionBuilder extends ConnectionBuilder {

    private RedisURI redisURI;

    public SslConnectionBuilder ssl(RedisURI redisURI) {
        this.redisURI = redisURI;
        return this;
    }

    public static SslConnectionBuilder sslConnectionBuilder() {
        return new SslConnectionBuilder();
    }

    @Override
    protected List<ChannelHandler> buildHandlers() {
        LettuceAssert.assertState(redisURI != null, "RedisURI must not be null");
        LettuceAssert.assertState(redisURI.isSsl(), "RedisURI is not configured for SSL (ssl is false)");

        return super.buildHandlers();
    }

    @Override
    @Deprecated
    public RedisChannelInitializer build() {

        return new SslChannelInitializer(getPingCommandSupplier(), this::buildHandlers, toHostAndPort(redisURI),
                redisURI.isVerifyPeer(), redisURI.isStartTls(), clientResources(), getTimeout(),
                clientOptions().getSslOptions());
    }

    @Override
    public RedisChannelInitializer build(SocketAddress socketAddress) {

        return new SslChannelInitializer(getPingCommandSupplier(), this::buildHandlers, toHostAndPort(socketAddress),
                redisURI.isVerifyPeer(), redisURI.isStartTls(), clientResources(), getTimeout(),
                clientOptions().getSslOptions());
    }

    static HostAndPort toHostAndPort(RedisURI redisURI) {

        if (LettuceStrings.isNotEmpty(redisURI.getHost())) {
            return HostAndPort.of(redisURI.getHost(), redisURI.getPort());
        }

        return null;
    }

    static HostAndPort toHostAndPort(SocketAddress socketAddress) {

        if (socketAddress instanceof InetSocketAddress) {

            InetSocketAddress isa = (InetSocketAddress) socketAddress;

            return HostAndPort.of(isa.getHostString(), isa.getPort());
        }

        return null;
    }

    /**
     * @author Mark Paluch
     */
    static class SslChannelInitializer extends io.netty.channel.ChannelInitializer<Channel> implements RedisChannelInitializer {

        private final Supplier<AsyncCommand<?, ?, ?>> pingCommandSupplier;
        private final Supplier<List<ChannelHandler>> handlers;
        private final HostAndPort hostAndPort;
        private final boolean verifyPeer;
        private final boolean startTls;
        private final ClientResources clientResources;
        private final Duration timeout;
        private final SslOptions sslOptions;

        private volatile CompletableFuture<Boolean> initializedFuture = new CompletableFuture<>();

        public SslChannelInitializer(Supplier<AsyncCommand<?, ?, ?>> pingCommandSupplier,
                Supplier<List<ChannelHandler>> handlers, HostAndPort hostAndPort, boolean verifyPeer, boolean startTls,
                ClientResources clientResources, Duration timeout, SslOptions sslOptions) {

            this.pingCommandSupplier = pingCommandSupplier;
            this.handlers = handlers;
            this.hostAndPort = hostAndPort;
            this.verifyPeer = verifyPeer;
            this.startTls = startTls;
            this.clientResources = clientResources;
            this.timeout = timeout;
            this.sslOptions = sslOptions;
        }

        @Override
        protected void initChannel(Channel channel) throws Exception {
            doInitialize(channel);
        }

        private void doInitialize(Channel channel) throws IOException, GeneralSecurityException {

            SSLEngine sslEngine = initializeSSLEngine(channel.alloc());

            if (channel.pipeline().get("first") == null) {
                channel.pipeline().addFirst("first", new ChannelDuplexHandler() {

                    @Override
                    public void channelActive(ChannelHandlerContext ctx) throws Exception {
                        clientResources.eventBus().publish(new ConnectedEvent(local(ctx), remote(ctx)));
                        super.channelActive(ctx);
                    }

                    @Override
                    public void channelInactive(ChannelHandlerContext ctx) throws Exception {
                        clientResources.eventBus().publish(new DisconnectedEvent(local(ctx), remote(ctx)));
                        super.channelInactive(ctx);
                    }
                });
            }

            SslHandler sslHandler = new SslHandler(sslEngine, startTls);
            channel.pipeline().addLast(sslHandler);

            if (channel.pipeline().get("channelActivator") == null) {
                channel.pipeline().addLast("channelActivator", new RedisChannelInitializerImpl() {

                    private AsyncCommand<?, ?, ?> pingCommand;

                    @Override
                    public CompletableFuture<Boolean> channelInitialized() {
                        return initializedFuture;
                    }

                    @Override
                    public void channelInactive(ChannelHandlerContext ctx) throws Exception {

                        if (!initializedFuture.isDone()) {
                            initializedFuture
                                    .completeExceptionally(new RedisConnectionException("Connection closed prematurely"));
                        }

                        initializedFuture = new CompletableFuture<>();
                        pingCommand = null;
                        super.channelInactive(ctx);
                    }

                    @Override
                    public void channelActive(ChannelHandlerContext ctx) throws Exception {
                        if (initializedFuture.isDone()) {
                            super.channelActive(ctx);
                        }
                    }

                    @Override
                    public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exception {
                        if (evt instanceof SslHandshakeCompletionEvent && !initializedFuture.isDone()) {

                            SslHandshakeCompletionEvent event = (SslHandshakeCompletionEvent) evt;
                            if (event.isSuccess()) {
                                if (pingCommandSupplier != PlainChannelInitializer.NO_PING) {
                                    pingCommand = pingCommandSupplier.get();
                                    pingBeforeActivate(pingCommand, initializedFuture, ctx, clientResources, timeout);
                                } else {
                                    ctx.fireChannelActive();
                                }
                            } else {
                                initializedFuture.completeExceptionally(event.cause());
                            }
                        }

                        if (evt instanceof ConnectionEvents.Activated) {
                            if (!initializedFuture.isDone()) {
                                initializedFuture.complete(true);
                                clientResources.eventBus().publish(new ConnectionActivatedEvent(local(ctx), remote(ctx)));
                            }
                        }

                        super.userEventTriggered(ctx, evt);
                    }

                    @Override
                    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
                        if (cause instanceof SSLHandshakeException || cause.getCause() instanceof SSLException) {
                            initializedFuture.completeExceptionally(cause);
                        }
                        super.exceptionCaught(ctx, cause);
                    }
                });
            }

            for (ChannelHandler handler : handlers.get()) {
                channel.pipeline().addLast(handler);
            }

            clientResources.nettyCustomizer().afterChannelInitialized(channel);
        }

        private SSLEngine initializeSSLEngine(ByteBufAllocator alloc) throws IOException, GeneralSecurityException {

            SSLParameters sslParams = sslOptions.createSSLParameters();
            SslContextBuilder sslContextBuilder = sslOptions.createSslContextBuilder();

            if (verifyPeer) {
                sslParams.setEndpointIdentificationAlgorithm("HTTPS");
            } else {
                sslContextBuilder.trustManager(InsecureTrustManagerFactory.INSTANCE);
            }

            SslContext sslContext = sslContextBuilder.build();

            SSLEngine sslEngine = hostAndPort != null
                    ? sslContext.newEngine(alloc, hostAndPort.getHostText(), hostAndPort.getPort())
                    : sslContext.newEngine(alloc);
            sslEngine.setSSLParameters(sslParams);

            return sslEngine;
        }

        @Override
        public CompletableFuture<Boolean> channelInitialized() {
            return initializedFuture;
        }
    }
}
